// cmd/onyx-snapshotter/main.go
package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	"github.com/containerd/containerd/v2/plugins"
	"github.com/onyx/snapshotter/pkg/config"
	"github.com/onyx/snapshotter/pkg/service"
	"google.golang.org/grpc"
)

func main() {
	var (
		root     = flag.String("root", "/var/lib/onyx-snapshotter", "root directory")
		address  = flag.String("address", "unix:///run/onyx-snapshotter/snapshotter.sock", "grpc listen address")
		logLevel = flag.String("log-level", "info", "log level")
	)
	flag.Parse()

	// 初始化配置
	cfg := &config.Config{
		Root:              *root,
		GRPCAddress:       *address,
		LogLevel:          *logLevel,
		ChunkSize:         4 * 1024 * 1024, // 4MB default
		DBBackend:         "rocksdb",       // or "boltdb"
		FanotifyEnabled:   true,
		UblkEnabled:       true,
		DeduplicationRate: 0.75, // target
	}

	// 确保目录存在
	for _, dir := range []string{
		cfg.Root,
		filepath.Join(cfg.Root, "chunks"),
		filepath.Join(cfg.Root, "metadata"),
		filepath.Join(cfg.Root, "snapshots"),
	} {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			log.Fatalf("failed to create directory %s: %v", dir, err)
		}
	}

	// 初始化服务
	snapshotter, err := service.NewSnapshotter(cfg)
	if err != nil {
		log.Fatalf("failed to create snapshotter: %v", err)
	}
	defer snapshotter.Close()

	// 启动 gRPC 服务
	if err := serveGRPC(cfg.GRPCAddress, snapshotter); err != nil {
		log.Fatalf("grpc server error: %v", err)
	}
}

func serveGRPC(addr string, snapshotter *service.Snapshotter) error {
	// 解析地址
	scheme, path := parseAddress(addr)
	var listener net.Listener
	var err error

	if scheme == "unix" {
		// 清理旧的 socket
		os.Remove(path)
		listener, err = net.Listen("unix", path)
		if err != nil {
			return fmt.Errorf("failed to listen on unix socket %s: %w", path, err)
		}
		// 设置权限
		os.Chmod(path, 0o660)
	} else if scheme == "tcp" {
		listener, err = net.Listen("tcp", path)
		if err != nil {
			return fmt.Errorf("failed to listen on tcp %s: %w", path, err)
		}
	} else {
		return fmt.Errorf("unsupported scheme: %s", scheme)
	}

	server := grpc.NewServer()
	plugins.Register(server, &plugins.Registration{
		Type: plugins.SnapshotPlugin,
		ID:   "onyx",
		InitFn: func(ic *plugins.InitContext) (interface{}, error) {
			return snapshotter, nil
		},
	})

	// 优雅关闭
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		<-sigCh
		log.Println("Shutting down gracefully...")
		server.GracefulStop()
	}()

	log.Printf("Starting onyx-snapshotter on %s", addr)
	return server.Serve(listener)
}

func parseAddress(addr string) (scheme, path string) {
	// 简单解析 "unix:///path" 或 "tcp://host:port"
	var i int
	for i = 0; i < len(addr); i++ {
		if addr[i] == ':' {
			break
		}
	}
	if i < len(addr) && addr[i+1:i+3] == "//" {
		scheme = addr[:i]
		path = addr[i+3:]
	}
	return scheme, path
}
